#include "yolt.h"
#include <ctime>
#include <cstring>
#include <opencv2/opencv.hpp>
#include <boost/filesystem.hpp>
// #include "../tools/files_io.h"

namespace nn_yolt
{

void NNYolt::construct_net()
{
    const float epsilon = 0.1;
    const int conv_size_3 = 3;
    const int conv_size_7 = 7;
    const int sub_sample_size_2 = 2;
    const int sub_sample_size_4 = 4;

    auto conv_type_same = tiny_dnn::padding::same;
    bool ct[1024 * 1024] = {false};
    auto make_table = [&](const int in_size, const int out_size) -> tiny_dnn::core::connection_table {
        gen_connection_table(in_size, out_size, ct);
        // for(int i = 0; i < in_size; i++)
        // {
        //     for(int j = 0; j < out_size; j++)
        //     {
        //         std::cout << ct[i*out_size + j] << " ";
        //     }
        //     std::cout << std::endl;
        // }
        return tiny_dnn::core::connection_table(ct, in_size, out_size);
    };

    net
        // C0: input: 416x416x3 ; output: 416x416x192
        << C(416, 416, conv_size_7, 3, 32, conv_type_same)
        << leaky_relu(epsilon)
        // S1: input: 416x416x192 ; output: 208x208x192
        << S(416, 416, 32, sub_sample_size_2)
        // C2: input: 208x208x192 ; output: 208x208x256
        << B(208 * 208, 32)
        << C(208, 208, conv_size_3, 32, 128, make_table(32, 128), conv_type_same)
        << leaky_relu(epsilon)
        << S(208, 208, 128, sub_sample_size_2)
        // C2: input: 208x208x192 ; output: 208x208x256
        << B(104 * 104, 128)
        << C(104, 104, conv_size_3, 128, 256, make_table(128, 256), conv_type_same)
        << leaky_relu(epsilon)
        // S3: input: 208x208x256 ; output: 104x104x256
        << S(104, 104, 256, sub_sample_size_2)
        // C4: input:104x104x256 ; output: 104x104x512
        // << B(104 * 104, 512)
        // << C(104, 104, conv_size_3, 256, 512, make_table(256, 512), conv_type_same)
        // << leaky_relu(epsilon)
        // S5: input: 104x104x512 ; output: 52x52x512
        // << S(104, 104, 512, sub_sample_size_2)
        // C6: input:52x52x512 ; output: 52x52x256
        << B(52 * 52, 256)
        << C(52, 52, conv_size_3, 256, 512, make_table(256, 512), conv_type_same)
        << leaky_relu(epsilon)
        << B(52 * 52, 512)
        << C(52, 52, conv_size_3, 512, 256, make_table(512, 256), conv_type_same)
        << leaky_relu(epsilon)
        // S7: input:52x52x256 ; output:26x26x256
        << S(52, 52, 256, sub_sample_size_2)
        // C8: input:26x26x256 ; output: 26x26x128
        << B(26 * 26, 256)
        << C(26, 26, conv_size_3, 256, 128, make_table(256, 128), conv_type_same)
        << leaky_relu(epsilon)
        // S9: input:26x26x128 ; output:13x13x128
        << S(26, 26, 128, sub_sample_size_2)
        // C10: input:13x13x128 ; output: 13x13x64
        << B(13 * 13, 128)
        << C(13, 13, conv_size_3, 128, 64, make_table(128, 64), conv_type_same)
        << leaky_relu(epsilon)
        // C11: input:13x13x64 ; output: 13x13x32
        << B(13 * 13, 64)
        << C(13, 13, conv_size_3, 64, 32, make_table(64, 32), conv_type_same)
        << leaky_relu(epsilon)
        // C12: input:13x13x32 ; output: 13x13x8
        << B(13 * 13, 32)
        << C(out_width, out_height, conv_size_3, 32, out_depth, conv_type_same)
        << relu();
        // F13: input:13x13x8 ; output: 13x13x20
        // << F(13 * 13 * 8, out_width * out_height * out_depth)
        // << leaky_relu(0.0f);

    // // BN
    // << B(52 * 52, 512)
    // // C8: input:52x52x512 ; output: 52x52x256
    // << C(52, 52, conv_size_3, 512, 256, make_table(512, 256), conv_type_same)
    // << leaky_relu(epsilon)
    // // S9: input: 52x52x256 ; output: 26x26x256
    // << S(52, 52, 256, sub_sample_size_2)
    // // << linear
    // // BN
    // << B(26 * 26, 256)
    // // C10: input:26x26x256 ; output: 26x26x512
    // << C(26, 26, conv_size_3, 256, 512, make_table(256, 512), conv_type_same)
    // << leaky_relu(epsilon)
    // // BN
    // // << B(26 * 26, 512)
    // // C11: input: 26x26x512 ; output: 26x26x10
    // << C(26, 26, conv_size_3, 256, 10, conv_type_same)
    // << leaky_relu(epsilon)
    // F8: input: 26x26x256 ; output: 4096
    // << F(26 * 26 * 256, 4096)
    // << leaky_relu(1.0f)
    // F9: input: 4096 ; output: 26x26x10
}

void NNYolt::gen_connection_table(const int in_size, const int out_size, bool *connection_table)
{
#define X false
#define O true

    int rows = in_size, cols = out_size;
    int non_zero_num = rows;

    // c[out_size * in_size]
    std::memset(connection_table, false, rows * cols * sizeof(bool));

    // 设置某一列
    auto set_col = [&](const int col, const int begin_row) -> void {
        for (int i = 0; i < non_zero_num; i++)
        {
            connection_table[((begin_row + i) % rows) * cols + col] = O;
        }
    };

    // 此后所有列，连接数减半
    int col = 1;
    // 第一列为 全连接
    set_col(0, 0);
    while (true)
    {
        non_zero_num /= 2;
        for (int row = 0; col < cols && row < rows; col++, row++)
        {
            set_col(col, row);
        }
        if (col == cols || non_zero_num == 0)
            return;
    }
#undef O
#undef X
}

void NNYolt::train(const str_t samples_dir, const int _begin_index,
                   const int train_sample_num, const int _batch_size, const int _epoch)
{
    tiny_dnn::adagrad opt;

    const int MAX_BATCH_SIZE = 20;
    const int MAX_ONCE_READ_NUM = 128;
    int batch_size = _batch_size > MAX_BATCH_SIZE ? MAX_BATCH_SIZE : _batch_size;
    int epoch = _epoch;
    int read_sample_num_each_round = train_sample_num > MAX_ONCE_READ_NUM ? MAX_ONCE_READ_NUM : train_sample_num;
    int round_num = train_sample_num / read_sample_num_each_round;
    int last_round_sample_num = train_sample_num % read_sample_num_each_round;

    int batch_index = 0, epoch_index = 0;
    char __tmp[64];
    time_t __curr_time;
    auto print_curr_time = [&]() -> void {
        std::time(&__curr_time);
        std::strftime(__tmp, sizeof(__tmp), "%Y-%m-%d %H:%M:%S", std::localtime(&__curr_time));

        std::cout << __tmp << " ";
    };
    auto each_batch = [&]() -> void {
        print_curr_time();
        std::cout << "Epoch: " << epoch_index << ", "
                  << "Batch: " << ++batch_index << std::endl;
    };
    auto each_epoch = [&]() -> void {
        print_curr_time();
        std::cout << "Epoch: " << epoch_index << std::endl;
        batch_index = 0;
    };

    int begin_index = 0;
    int end_index = _begin_index + train_sample_num;
    std::vector<f_vec_t> in_data, out_data;
    auto train_once_round = [&]() -> void {
        in_data.clear();
        out_data.clear();
        get_train_data(samples_dir, in_data, out_data, begin_index, read_sample_num_each_round);
        net.fit<tiny_dnn::yolt_error>(opt, in_data, out_data, batch_size, 1, each_batch, []() { ; });
        // net.save(net_file_name);
    };

    // 每一次循环为一轮训练
    for (epoch_index = 1; epoch_index <= epoch; epoch_index++)
    {
        each_epoch();

        // 每一轮训练，分批次读取数据
        for (begin_index = _begin_index; begin_index < end_index; begin_index += read_sample_num_each_round)
        {
            in_data.clear();
            out_data.clear();
            get_train_data(samples_dir, in_data, out_data, begin_index, read_sample_num_each_round);
            net.fit<tiny_dnn::yolt_error>(opt, in_data, out_data, batch_size, 1, each_batch, []() { ; });
        }
        // 读取剩余数据
        if (last_round_sample_num != 0)
        {
            in_data.clear();
            out_data.clear();
            get_train_data(samples_dir, in_data, out_data, begin_index, read_sample_num_each_round);
            net.fit<tiny_dnn::yolt_error>(opt, in_data, out_data, batch_size, 1, each_batch, []() { ; });
        }
    }
}

void NNYolt::get_train_data(const str_t samples_dir, std::vector<f_vec_t> &in_vec,
                            std::vector<f_vec_t> &out_vec,
                            const int train_sample_index_begin,
                            const int train_sample_num) const
{
    assert(boost::filesystem::exists(samples_dir));
    assert(boost::filesystem::is_directory(samples_dir));

    const str_t label_suffix = ".txt";

    s_vec_t image_paths, label_paths;
    boost::filesystem::directory_iterator iter(samples_dir);
    boost::filesystem::directory_iterator iter_end;
    for (int i = 0; iter != iter_end; iter++)
    {
        const auto &curr_file_path = iter->path();

        if (boost::filesystem::is_regular_file(curr_file_path) && curr_file_path.extension() == ".png")
        {
            str_t str = curr_file_path.string();
            image_paths.push_back(str);
            label_paths.push_back(str.replace(str.rfind('.'), label_suffix.size(), label_suffix));
            i++;
        }
    }

    std::sort(image_paths.begin(), image_paths.end());
    std::sort(label_paths.begin(), label_paths.end());

    // 将图片和 txt 数据转换为 yolt 网络可接受的格式
    int n = train_sample_num <= 0 ? image_paths.size() : (train_sample_num + train_sample_index_begin);
    n = n > image_paths.size() ? image_paths.size() : n;
    YoltLabel curr_label(in_width, in_height, out_width, out_height, box_num_per_cell);
    for (int i = train_sample_index_begin; i < n; i++)
    {
        f_vec_t in, out;

        // std::cout << "Current image name: " << image_paths[i] << std::endl;
        // std::cout << "Current txt name: " << label_paths[i] << std::endl;

        curr_label.load(label_paths[i]);
        mat_2_in_data(image_paths[i], in);
        label_2_out_data(curr_label, out);

        // std::cout << net.in_data_size() << std::endl;
        // std::cout << in.size() << std::endl;
        // std::cout << out.size() << std::endl;

        in_vec.push_back(in);
        out_vec.push_back(out);
    }
}

void NNYolt::mat_2_in_data(const str_t image_path, f_vec_t &vec) const
{
    cv::Mat image = cv::imread(image_path, cv::IMREAD_COLOR);
    assert(image.data);

    for (int i = 0; i < image.rows; i++)
    {
        for (int j = 0; j < image.cols; j++)
        {
            const cv::Vec3b &v = image.at<cv::Vec3b>(i, j);
            for (int k = 0; k < image.channels(); k++)
            {
                vec.push_back(v[k]);
            }
        }
    }
}

void NNYolt::label_2_out_data(const YoltLabel &label, f_vec_t &vec) const
{
    for (int i = 0; i < out_width; i++)
    {
        for (int j = 0; j < out_height; j++)
        {
            for (int k = 0; k < box_num_per_cell; k++)
            {
                f_vec_t params = label.get(i, j, k);
                for (int p = 0; p < PARAM_NUM_PER_BOX; p++)
                {
                    vec.push_back(params[p]);
                }
            }
        }
    }
}

void NNYolt::in_data_2_mat(const f_vec_t &in_data, cv::Mat &mat) const
{
    assert(mat.data);
    assert(in_data.size() == mat.rows * mat.cols * mat.channels());

    int index = 0;
    for (int i = 0; i < mat.rows; i++)
    {
        for (int j = 0; j < mat.cols; j++)
        {
            cv::Vec3b &v = mat.at<cv::Vec3b>(i, j);
            for (int k = 0; k < mat.channels(); k++)
            {
                v[k] = in_data[index++];
            }
        }
    }
}

void NNYolt::out_data_2_label(const f_vec_t &out_data, YoltLabel &label) const
{
    assert(out_data.size() == out_width * out_height * out_depth);

    int index = 0;
    f_vec_t params(PARAM_NUM_PER_BOX);

    for (int i = 0; i < out_width; i++)
    {
        for (int j = 0; j < out_height; j++)
        {
            for (int k = 0; k < box_num_per_cell; k++)
            {
                for (int p = 0; p < PARAM_NUM_PER_BOX; p++)
                {
                    params[p] = out_data[index++];
                }
                label.set(i, j, k, params);
            }
        }
    }
}

void NNYolt::test(const str_t samples_dir, const int test_index_begin, const int test_num)
{
    std::vector<f_vec_t> in_data, desired_data;

    get_train_data(samples_dir, in_data, desired_data, test_index_begin, test_num);

    std::vector<f_vec_t> out_data = net.test(in_data);

    int cell_width = in_width / out_width;
    int cell_height = in_height / out_height;
    const float_t iou_lower = 0.5f;

    for (int i = 0; i < test_num; i++)
    {
        cv::Mat image = cv::Mat::zeros(in_width, in_height, CV_8UC3);
        YoltLabel label(in_width, in_height, out_width, out_height, box_num_per_cell);
        cv::Scalar color(0, 0, 255);

        in_data_2_mat(in_data[i], image);
        out_data_2_label(out_data[i], label);
        label.denormalize();
        for (int cell_ny = 0; cell_ny < out_height; cell_ny++)
        {
            for (int cell_nx = 0; cell_nx < out_width; cell_nx++)
            {
                for (int box = 0; box < box_num_per_cell; box++)
                {
                    f_vec_t params = label.get(cell_nx, cell_ny, box);
                    std::cout << params[4] << std::endl;
                    if (params[4] >= iou_lower) // iou != 0
                    {
                        int cx = cell_width * cell_nx + params[0];
                        int cy = cell_height * cell_ny + params[1];
                        cv::rectangle(image, cv::Rect(cx - params[2] / 2, cy - params[3] / 2, params[2], params[3]), color, 2);
                    }
                }
            }
        }

        cv::imshow("Image", image);
        cv::waitKey(0);
    }
}

} // namespace nn_yolt